import torch
import torch.nn as nn
import torch.optim as optim
import random
import networkx as nx
import tkinter as tk
import time
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 检查 GPU

# 构建 DAG 并赋予随机权重
def build_dag_with_weights(layers):
    G = nx.DiGraph()
    prev_layer_nodes = [0]
    node_counter = 1
    weights = {}

    for layer_size in layers:
        current_layer_nodes = list(range(node_counter, node_counter + layer_size))
        node_counter += layer_size

        if len(prev_layer_nodes) == 1 and layer_size > 1:
            for node in current_layer_nodes:
                G.add_edge(prev_layer_nodes[0], node)
                weights[(prev_layer_nodes[0], node)] = random.uniform(1, 10)
        elif len(prev_layer_nodes) > 1 and layer_size == 1:
            for node in prev_layer_nodes:
                G.add_edge(node, current_layer_nodes[0])
                weights[(node, current_layer_nodes[0])] = random.uniform(1, 10)
        else:
            for prev_node in prev_layer_nodes:
                for curr_node in current_layer_nodes:
                    G.add_edge(prev_node, curr_node)
                    weights[(prev_node, curr_node)] = random.uniform(1, 10)

        prev_layer_nodes = current_layer_nodes

    return G, weights


# DQN 模型
class DQN(nn.Module):
    def __init__(self, state_size, action_size):
        super(DQN, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(state_size, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, action_size),
        )

    def forward(self, x):
        return self.fc(x)


# 训练 DQN
class DQNAgent:
    def __init__(self, state_size, action_size, gamma=0.99, lr=0.001, epsilon=1.0, epsilon_decay=0.995, epsilon_min=0.01):
        self.state_size = state_size
        self.action_size = action_size
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min

        self.model = DQN(state_size, action_size).to(device)
        self.optimizer = optim.Adam(self.model.parameters(), lr=lr)
        self.criterion = nn.MSELoss()

        self.memory = []

    def act(self, state):
        if random.random() < self.epsilon:
            return random.randint(0, self.action_size - 1)
        state = torch.FloatTensor(state).unsqueeze(0).to(device)
        with torch.no_grad():
            q_values = self.model(state)
        return torch.argmax(q_values).item()

    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))
        if len(self.memory) > 1000:
            self.memory.pop(0)

    def replay(self, batch_size=32):
        if len(self.memory) < batch_size:
            return

        batch = random.sample(self.memory, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states = torch.FloatTensor(states).to(device)
        actions = torch.LongTensor(actions).to(device)
        rewards = torch.FloatTensor(rewards).to(device)
        next_states = torch.FloatTensor(next_states).to(device)
        dones = torch.FloatTensor(dones).to(device)

        # 计算 Q 值
        q_values = self.model(states).gather(1, actions.unsqueeze(-1)).squeeze(-1)

        # 计算目标 Q 值
        with torch.no_grad():
            next_q_values = self.model(next_states).max(1)[0]
            targets = rewards + self.gamma * next_q_values * (1 - dones)

        # 更新模型
        self.optimizer.zero_grad()
        loss = self.criterion(q_values, targets)
        loss.backward()
        self.optimizer.step()

        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay


# 可视化路径选择
def visualize_path(graph, weights, path, canvas, node_positions):
    canvas.delete("all")
    # 绘制边
    for edge, weight in weights.items():
        x1, y1 = node_positions[edge[0]]
        x2, y2 = node_positions[edge[1]]
        edge_color = "gray"  # 默认颜色
        if edge in zip(path[:-1], path[1:]):  # 高亮路径中的边
            edge_color = "red"
        canvas.create_line(x1, y1, x2, y2, fill=edge_color, width=3 if edge_color == "red" else 2)
        canvas.create_text((x1 + x2) / 2, (y1 + y2) / 2, text=f"{weight:.1f}", fill="black")

    # 绘制节点
    for node in graph.nodes:
        x, y = node_positions[node]
        color = "blue"  # 默认颜色
        if node in path:  # 高亮路径中的节点
            color = "red"
        canvas.create_oval(x - 10, y - 10, x + 10, y + 10, fill=color)

    canvas.update()
    time.sleep(0.5)  # 动态展示路径选择过程


# 主程序
def main():
    layers = [1, 1, 4, 1, 4, 1, 3, 1]
    G, weights = build_dag_with_weights(layers)

    state_size = len(G.nodes)
    action_size = max(len(list(G.successors(n))) for n in G.nodes)

    agent = DQNAgent(state_size, action_size)

    # Tkinter 可视化
    root = tk.Tk()
    root.title("DQN Pathfinding in DAG")
    canvas = tk.Canvas(root, width=800, height=600, bg="white")
    canvas.pack()

    node_positions = nx.spring_layout(G, seed=42)
    node_positions = {k: (int(400 + v[0] * 300), int(300 + v[1] * 300)) for k, v in node_positions.items()}

    # 训练 DQN
    for episode in range(200):
        state = [0] * state_size
        current_node = 0
        state[current_node] = 1
        total_reward = 0
        path = [current_node]

        while True:
            successors = list(G.successors(current_node))
            action = agent.act(state)
            if action >= len(successors):
                action = random.randint(0, len(successors) - 1)
            next_node = successors[action]

            reward = weights[(current_node, next_node)]
            next_state = [0] * state_size
            next_state[next_node] = 1

            done = next_node == max(G.nodes)
            agent.remember(state, action, reward, next_state, done)
            state = next_state
            current_node = next_node
            total_reward += reward
            path.append(current_node)

            visualize_path(G, weights, path, canvas, node_positions)

            if done:
                break

        agent.replay(batch_size=64)  # 增加批量训练样本，提升 GPU 性能
        print(f"Episode {episode + 1}: Total reward = {total_reward:.2f}")

    root.mainloop()


if __name__ == "__main__":
    main()
